"""
优先级评估模块（决策树模型）
功能：结合历史处理记录评估事件优先级
优先级等级：紧急、一般、可延后
输入：事件分类结果 + 历史处理数据
输出：优先级标签
"""

import pandas as pd
from typing import List, Dict, Tuple
from sklearn.tree import DecisionTreeClassifier
from sklearn.preprocessing import LabelEncoder

class PriorityEvaluator:
    def __init__(self):
        # 优先级标签
        self.priority_levels = ['紧急', '一般', '可延后']
        # 决策树分类器
        self.model = DecisionTreeClassifier(random_state=42)
        # 标签编码器
        self.label_encoders = {}
        # 是否已训练
        self.is_trained = False
    
    def _encode_features(self, data: List[Dict]) -> pd.DataFrame:
        """
        对特征进行编码
        :param data: 原始数据列表
        :return: 编码后的DataFrame
        """
        df = pd.DataFrame(data)
        
        # 对分类特征进行标签编码
        categorical_features = ['event_type', 'location_type', 'time_period']
        for feature in categorical_features:
            if feature in df.columns:
                if feature not in self.label_encoders:
                    self.label_encoders[feature] = LabelEncoder()
                    df[feature] = self.label_encoders[feature].fit_transform(df[feature])
                else:
                    df[feature] = self.label_encoders[feature].transform(df[feature])
        
        return df
    
    def train(self, training_data: List[Dict], labels: List[str]) -> bool:
        """
        训练优先级评估模型
        :param training_data: 训练数据列表
        :param labels: 对应优先级标签列表
        :return: 是否训练成功
        """
        try:
            # 验证标签
            for label in labels:
                if label not in self.priority_levels:
                    raise ValueError(f"无效优先级标签: {label}")
            
            # 编码特征
            df = self._encode_features(training_data)
            
            # 分离特征和标签
            X = df.drop(columns=['id'], errors='ignore')
            y = labels
            
            # 训练模型
            self.model.fit(X, y)
            self.is_trained = True
            return True
        except Exception as e:
            print(f"模型训练失败: {e}")
            return False
    
    def evaluate(self, event_data: Dict) -> Tuple[str, float]:
        """
        评估单个事件的优先级
        :param event_data: 事件数据
        :return: (优先级标签, 置信度)
        """
        if not self.is_trained:
            raise RuntimeError("模型未训练，请先调用train方法")
            
        # 编码特征
        df = self._encode_features([event_data])
        X = df.drop(columns=['id'], errors='ignore')
        
        # 预测
        prediction = self.model.predict(X)[0]
        probabilities = self.model.predict_proba(X)[0]
        confidence = max(probabilities)
        
        return prediction, confidence
    
    def evaluate_batch(self, events_data: List[Dict]) -> List[Tuple[str, float]]:
        """
        批量评估事件优先级
        :param events_data: 事件数据列表
        :return: [(优先级标签, 置信度), ...]
        """
        if not self.is_trained:
            raise RuntimeError("模型未训练，请先调用train方法")
            
        # 编码特征
        df = self._encode_features(events_data)
        X = df.drop(columns=['id'], errors='ignore')
        
        # 预测
        predictions = self.model.predict(X)
        probabilities = self.model.predict_proba(X)
        
        results = []
        for i, pred in enumerate(predictions):
            confidence = max(probabilities[i])
            results.append((pred, confidence))
            
        return results
    
    def get_priority_levels(self) -> List[str]:
        """
        获取优先级标签列表
        :return: 优先级标签列表
        """
        return self.priority_levels

# 使用示例
def main():
    evaluator = PriorityEvaluator()
    
    # 训练数据
    training_data = [
        {'id': 'EV001', 'event_type': '交通事故', 'location_type': '主干道', 'time_period': '高峰时段', 'impact_scope': 5},
        {'id': 'EV002', 'event_type': '施工', 'location_type': '次干道', 'time_period': '平峰时段', 'impact_scope': 3},
        {'id': 'EV003', 'event_type': '设施故障', 'location_type': '地铁站', 'time_period': '高峰时段', 'impact_scope': 4},
        {'id': 'EV004', 'event_type': '拥堵', 'location_type': '主干道', 'time_period': '高峰时段', 'impact_scope': 5},
        {'id': 'EV005', 'event_type': '其他', 'location_type': '社区道路', 'time_period': '夜间', 'impact_scope': 2},
        {'id': 'EV006', 'event_type': '交通事故', 'location_type': '高速公路', 'time_period': '平峰时段', 'impact_scope': 4},
        {'id': 'EV007', 'event_type': '施工', 'location_type': '主干道', 'time_period': '夜间', 'impact_scope': 4},
        {'id': 'EV008', 'event_type': '设施故障', 'location_type': '公交站', 'time_period': '平峰时段', 'impact_scope': 3}
    ]
    
    training_labels = ['紧急', '一般', '紧急', '紧急', '可延后', '紧急', '一般', '一般']
    
    # 训练模型
    if evaluator.train(training_data, training_labels):
        print("优先级评估模型训练成功")
        
        # 测试数据
        test_data = [
            {'id': 'EV101', 'event_type': '交通事故', 'location_type': '主干道', 'time_period': '高峰时段', 'impact_scope': 5},
            {'id': 'EV102', 'event_type': '施工', 'location_type': '社区道路', 'time_period': '夜间', 'impact_scope': 2}
        ]
        
        # 批量评估
        results = evaluator.evaluate_batch(test_data)
        
        print("\n优先级评估结果:")
        for i, (priority, confidence) in enumerate(results):
            print(f"  事件 {test_data[i]['id']}: {priority} (置信度: {confidence:.2f})")
            print(f"    类型: {test_data[i]['event_type']}, 位置: {test_data[i]['location_type']}, 时间: {test_data[i]['time_period']}")
    else:
        print("优先级评估模型训练失败")

if __name__ == "__main__":
    main()